"""
Interface definitions for classification heads in continual learning.

Classification heads convert feature representations from backbone networks
into task-specific outputs. This module defines interfaces ensuring
consistent behavior across different head implementations.
"""

from abc import abstractmethod
from typing import Optional

import torch
from torch import nn

from continuallearning.interfaces.core.component import ComponentInterface
from continuallearning.interfaces.core.task_related_interface import (
    TaskIrrelevantInterface,
)
from typing import Dict, Any, List


class HeadInterface(ComponentInterface, TaskIrrelevantInterface):
    """
    Interface for model head implementations.

    Classification heads take feature representations and transform them
    into task-specific outputs (typically class logits). In continual learning,
    heads can be shared across tasks or task-specific.
    """

    @abstractmethod
    def forward(
        self,
        features: torch.Tensor,
        task_ids: Optional[List[int]] = None,
        *args,
        **kwargs,
    ) -> torch.Tensor:
        """
        Process features to produce outputs.

        Args:
            features: Input features from backbone of shape [B, in_features]
            task_ids: Optional list of task identifiers for multi-head models

        Returns:
            torch.Tensor: Model outputs (e.g., logits) of shape [B, num_classes]
        """
        pass


class MultiHeadInterface(nn.Module, HeadInterface):
    """
    Protocol for multi-head classification heads.

    Multi-head models maintain separate classifiers per task.
    This protocol extends HeadInterface with task-specific methods.
    """

    @abstractmethod
    def prepare_task(self, task_id: int, **kwargs) -> None:
        pass

    @abstractmethod
    def register_task(self, task_id: int, **kwargs) -> None:
        """
        Register a new task with this hook.

        Args:
            task_id: Identifier for the new task
            config: Optional task-specific configuration
        """
        pass

    @abstractmethod
    def set_task(self, task_id: int) -> None:
        """
        Configure the model for a specific task.

        Args:
            task_id: Task identifier
        """
        self._current_task = task_id

    @abstractmethod
    def set_trainable_tasks(self, task_ids: int | List[int], **kwargs) -> None:
        """
        Prepare this hook for a specific task.

        Args:
            task_id: Identifier for the task to prepare
            **kwargs: Additional preparation parameters
        """
        pass

    @abstractmethod
    def _freeze(self, task_ids: int | List[int]) -> None:
        """
        Freeze parameters for a specific task.

        Args:
            task_id: Task identifier to freeze parameters for
        """
        pass

    @abstractmethod
    def _unfreeze(self, task_ids: int | List[int]) -> None:
        """
        Freeze parameters for a specific task.

        Args:
            task_id: Task identifier to freeze parameters for
        """
        pass

    @property
    @abstractmethod
    def current_task(self) -> int:
        """Get the current task ID."""
        pass

    @property
    @abstractmethod
    def registered_tasks(self) -> List[int]:
        """Get list of registered task IDs."""
        pass

    @abstractmethod
    def process_with_task(
        self, x: torch.Tensor, task_id: int, **kwargs
    ) -> torch.Tensor:
        """Process input with a specific task."""
        pass

    @abstractmethod
    def process_with_tasks(
        self, x: torch.Tensor, task_ids: Optional[List[int]] = None, **kwargs
    ) -> torch.Tensor:
        """Process input among multiple tasks."""
        pass
